import logging
from abc import ABC, abstractmethod

from neo4j import GraphDatabase
from sqlalchemy import create_engine, inspect

logger = logging.getLogger(__name__)


# -----------------------------------------------------------------------------
# --- 抽象基类 (Base Class) - 遵循建模、提取、转换、加载思想 ---
# -----------------------------------------------------------------------------
class BaseMigrator(ABC):
    """
    定义了从 RDBMS 到 Neo4j 迁移的 M-ETL (Model, Extract, Transform, Load) 流程框架。
    这个基类只负责“做什么”(What)，具体的“怎么做”(How)由子类实现。
    """

    def __init__(self, rdbms_config, neo4j_config, batch_size=1000):
        rdb_url = f"{rdbms_config['dialect']}+{rdbms_config['driver']}://{rdbms_config['user']}:{rdbms_config['password']}@{rdbms_config['host']}:{rdbms_config['port']}/{rdbms_config['database']}"
        self.rdbms_engine = create_engine(rdb_url)
        self.inspector = inspect(self.rdbms_engine)
        self.neo4j_driver = GraphDatabase.driver(
            neo4j_config['uri'], auth=(neo4j_config['user'], neo4j_config['password'])
        )
        self.graph_model = None  # 模型将在 model_schema() 中被创建
        self.batch_size = batch_size
        logging.info("BaseMigrator initialized.")

    def close(self):
        self.neo4j_driver.close()
        logging.info("Connections closed.")

    # --- 抽象方法 (子类必须实现) ---
    @abstractmethod
    def model_schema(self):
        """
        步骤1: 数据建模 (Model)。
        子类必须实现此方法，以填充 self.graph_model。
        例如，通过检查外键，或调用 LLM。
        """
        raise NotImplementedError

    @abstractmethod
    def transform_data(self, data_row, context):
        """
        步骤3: 数据转换 (Transform)。
        子类必须实现此方法，以将提取的原始数据转换为加载所需的目标格式。
        """
        raise NotImplementedError

    # --- 具体方法 (由基类提供) ---
    def extract_data(self, source_name):
        """步骤2: 数据提取 (Extract)。这是一个通用实现。"""
        logging.info(f"[EXTRACT] Extracting data from source: {source_name}")
        query = f"SELECT * FROM {source_name}"
        with self.rdbms_engine.connect() as connection:
            from sqlalchemy.sql import text
            result = connection.execute(text(query))
            for row in result:
                yield dict(row._mapping)

    def load_data(self, batch, context):
        """步骤4: 数据加载 (Load)。这是一个通用实现。"""
        if not batch: return

        if context['type'] == 'node':
            with self.neo4j_driver.session() as session:
                query = f"UNWIND $batch as properties MERGE (n:{context['label']} {{{context['pk']}: properties.{context['pk']}}}) SET n = properties"
                session.run(query, batch=batch)
        elif context['type'] == 'relationship':
            with self.neo4j_driver.session() as session:
                query = f"""
                UNWIND $batch AS pair
                MATCH (start:{context['start_node_label']} {{{context['start_node_key']}: pair.start_id}})
                MATCH (end:{context['end_node_label']} {{{context['end_node_key']}: pair.end_id}})
                MERGE (start)-[:{context['rel_type']}]->(end)
                """
                session.run(query, batch=batch)

    # --- 流程编排 (Orchestration) ---
    def run_migration(self):
        """执行完整的 M-ETL 迁移流程。"""
        try:
            # 1. 建模
            self.model_schema()
            if not self.graph_model:
                raise ValueError("Graph model was not created. Aborting migration.")

            # 准备工作: 创建约束
            logging.info("Preparing Neo4j: creating constraints...")
            with self.neo4j_driver.session() as session:
                for node_model in self.graph_model['nodes'].values():
                    if node_model['primary_key']:
                        session.run(
                            f"CREATE CONSTRAINT IF NOT EXISTS FOR (n:{node_model['label_name']}) REQUIRE n.{node_model['primary_key']} IS UNIQUE")
            logging.info("Constraints are set.")

            # 迁移节点 (E-T-L)
            for table_name, node_model in self.graph_model['nodes'].items():
                if node_model.get("is_relationship_table"): continue  # 跳过被标记为纯关系表的节点
                batch = []
                context = {"type": "node", "label": node_model['label_name'], "pk": node_model['primary_key']}
                for row in self.extract_data(table_name):
                    transformed_row = self.transform_data(row, context)
                    if transformed_row:
                        batch.append(transformed_row)
                    if len(batch) >= self.batch_size:
                        self.load_data(batch, context)
                        batch = []
                self.load_data(batch, context)

            # 迁移关系 (E-T-L)
            for rel_model in self.graph_model['relationships']:
                batch = []
                context = {"type": "relationship", **rel_model}
                for row in self.extract_data(rel_model['source_table_name']):
                    transformed_row = self.transform_data(row, context)
                    if transformed_row:
                        batch.append(transformed_row)
                    if len(batch) >= self.batch_size:
                        self.load_data(batch, context)
                        batch = []
                self.load_data(batch, context)

            logging.info("Migration completed successfully!")
        except Exception as e:
            logging.error(f"An error occurred during migration: {e}", exc_info=True)
        finally:
            self.close()
